package org.example.fboot.config;

import io.jsonwebtoken.Claims;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.example.fboot.MessageHandler;
import org.example.fboot.handler.BusException;
import org.example.fboot.jwt.JwtUtils;
import org.example.fboot.proto.MessageRequestOuterClass;
import org.example.fboot.proto.RequestMessageOuterClass;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.util.StringUtils;
import org.springframework.web.reactive.HandlerMapping;
import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
import org.springframework.web.reactive.socket.HandshakeInfo;
import org.springframework.web.reactive.socket.WebSocketHandler;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

@Slf4j
@Configuration
public class WebSocketConfig {
    @Resource
    private MessageHandler<MessageRequestOuterClass.MessageRequest> handler;
    @Resource
    private WsPostProcessor postProcessor;

    @Bean
    public HandlerMapping webSocketHandlerMapping() {
        log.info("webSocketHandlerMapping init");
        Map<String, WebSocketHandler> urlMap = new HashMap<>();
        urlMap.put(handler.path(), (session) -> {
            HandshakeInfo info = session.getHandshakeInfo();
            HttpHeaders headers = info.getHeaders();
            String jwt = headers.getFirst(HttpHeaders.AUTHORIZATION);
            jwt = StringUtils.hasText(jwt) ? jwt : headers.getFirst("Sec-Websocket-Protocol");
            AtomicReference<String> wsId = new AtomicReference<>();
            Claims claims = JwtUtils.parseClaim(jwt).getPayload();
            wsId.set((String) claims.get("wsId"));
            return session.receive()
                    .doOnSubscribe((subscription) -> handler.session(wsId.get(), session))
                    .doOnTerminate(() -> handler.out(wsId.get(), session))
                    .doOnNext(webSocketMessage -> {
                        try {
                            //MessageRequestOuterClass.MessageRequest request = MessageRequestOuterClass.MessageRequest.parseFrom(webSocketMessage.getPayload().asInputStream());
                            RequestMessageOuterClass.RequestMessage requestMessage = RequestMessageOuterClass.RequestMessage.parseFrom(webSocketMessage.getPayload().asInputStream());
                            postProcessor.precess(requestMessage,(result)->{
                                //返回结果
                                session.send(Mono.fromCallable(()->session.binaryMessage(b->b.wrap(result.toByteArray())))).subscribe();
//                                handler.sender(request);
                            });
                        } catch (IOException | InvocationTargetException | IllegalAccessException e) {
                            throw new RuntimeException(e);
                        }
                    }).onErrorMap(e -> {
                        log.error("发送消息异常", e);
                        handler.out(wsId.get(), session);
                        return new BusException(e.getMessage());
                    }).then();
        });
        SimpleUrlHandlerMapping handlerMapping = new SimpleUrlHandlerMapping();
        handlerMapping.setOrder(1);
        handlerMapping.setUrlMap(urlMap);
        return handlerMapping;
    }
}
